'''
Author: Ryan
Email: 409266255@qq.com
Date: 2025-09-23 22:11:45

LastEditors: Ryan
LastEditTime: 2025-09-24 19:40:55
Description: tushare数据获取

FilePath: /qopilot/qopilot/data/datasource/qop_tushare.py
Copyright (c) 2025 by Ryan, All Rights Reserved. 
'''
import tushare as ts
import pandas as pd
from typing import Optional
import logging
from enum import Enum

from qopilot.config import config

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class HKConstants(Enum):
    """港股相关常量枚举"""

    # 港股列表字段
    HK_BASIC_FIELDS = [
        "ts_code",      # 股票代码
        "name",         # 股票简称
        "fullname",     # 公司全称
        "enname",       # 英文名称
        "cn_spell",     # 拼音
        "market",       # 市场类别
        "list_status",  # 上市状态
        "list_date",    # 上市日期
        "delist_date",  # 退市日期
        "trade_unit",   # 交易单位
        "isin",         # ISIN代码
        "curr_type",    # 货币代码
    ]

    # 港股交易日历字段
    HK_TRADE_CAL_FIELDS = [
        "cal_date",         # 日历日期
        "is_open",          # 是否交易 '0'休市 '1'交易
        # "pretrade_date",    # 上一个交易日
    ]

    # 港股复权行情字段
    HK_DAILY_ADJ_FIELDS = [
        "ts_code",          # 股票代码
        "trade_date",       # 交易日期
        "close",            # 收盘价
        "open",             # 开盘价
        "high",             # 最高价
        "low",              # 最低价
        # "pre_close",        # 昨收价
        # "change",           # 涨跌额
        # "pct_change",       # 涨跌幅
        "vol",              # 成交量
        "amount",           # 成交额
        # "vwap",             # 平均价
        # "adj_factor",       # 复权因子
        # "turnover_ratio",   # 换手率(基于总股本)
        # "free_share",       # 流通股本
        # "total_share",      # 总股本
        # "free_mv",          # 流通市值
        # "total_mv",         # 总市值
    ]

    # 市场状态常量
    LIST_STATUS_ACTIVE = "L"        # 上市
    LIST_STATUS_DELISTED = "D"      # 退市
    LIST_STATUS_PAUSED = "P"        # 暂停上市

    # 交易状态常量
    TRADE_OPEN = "1"                # 交易
    TRADE_CLOSED = "0"              # 休市


class HKTushareData:
    """港股Tushare数据获取类"""

    def __init__(self, token: Optional[str] = None):
        """
        初始化港股数据获取器

        Args:
            token: Tushare token，如果不提供则从配置文件读取
        """
        self.token = token or config.TUSHARE_TOKEN
        if not self.token:
            raise ValueError("请提供Tushare token或设置TUSHARE_TOKEN环境变量")

        # 初始化Tushare
        ts.set_token(self.token)
        self.pro = ts.pro_api()

        logger.info("港股数据获取器初始化成功")

    def get_hk_basic(self, 
                     ts_code: Optional[str] = None,
                     list_status: Optional[str] = None,
                     fields: Optional[list] = HKConstants.HK_BASIC_FIELDS.value) -> pd.DataFrame:
        """
        获取港股列表

        Args:
            ts_code: 股票代码，默认为None
            list_status: 上市状态，默认为None
            fields: 返回字段，默认为HKConstants.HK_BASIC_FIELDS.value

        Returns:
            DataFrame: 港股列表数据（已过滤掉ts_code以'8'开头的数据）

        Api Ref:
            https://tushare.pro/document/2?doc_id=191
        """
        try:
            logger.info("开始获取港股列表")

            # 调用Tushare接口
            df = self.pro.hk_basic(
                ts_code=ts_code,
                list_status=list_status,
                fields=fields
            )

            # 过滤掉ts_code以'8'开头的数据
            if not df.empty and 'ts_code' in df.columns:
                original_count = len(df)
                df = df[~df['ts_code'].str.startswith('8')]
                filtered_count = len(df)
                logger.info(f"过滤前共{original_count}条记录，过滤后剩余{filtered_count}条记录")
            else:
                logger.info("数据为空或缺少ts_code字段，无需过滤")

            logger.info(f"成功获取港股列表，共{len(df)}条记录")
            return df

        except Exception as e:
            logger.error(f"获取港股列表失败: {str(e)}")
            raise

    def get_hk_tradecal(self, 
                        start_date: Optional[str] = None,
                        end_date: Optional[str] = None,
                        is_open: Optional[str] = None,
                        fields: Optional[list] = HKConstants.HK_TRADE_CAL_FIELDS.value) -> pd.DataFrame:
        """
        获取港股交易日历

        Args:
            start_date: 开始日期，格式为'YYYYMMDD'，默认为None
            end_date: 结束日期，格式为'YYYYMMDD'，默认为None
            is_open: 是否交易，'0'表示休市，'1'表示交易，默认为None
            fields: 返回字段，默为HKConstants.HK_TRADE_CAL_FIELDS.value

        Returns:
            DataFrame: 港股交易日历数据

        Api Ref:
            https://tushare.pro/document/2?doc_id=250
        """
        try:
            logger.info("开始获取港股交易日历")

            # 调用Tushare接口
            df = self.pro.hk_tradecal(
                start_date=start_date,
                end_date=end_date,
                is_open=is_open,
                fields=fields
            )

            logger.info(f"成功获取港股交易日历，共{len(df)}条记录")
            return df

        except Exception as e:
            logger.error(f"获取港股交易日历失败: {str(e)}")
            raise

    def get_hk_daily_adj(self, 
                         ts_code: Optional[str] = None,
                         trade_date: Optional[str] = None,
                         start_date: Optional[str] = None,
                         end_date: Optional[str] = None,
                         fields: Optional[list] = HKConstants.HK_DAILY_ADJ_FIELDS.value) -> pd.DataFrame:
        """
        获取港股复权行情数据

        Args:
            ts_code: 股票代码，格式为'00001.HK'，默认为None
            trade_date: 交易日期，格式为'YYYYMMDD'，默认为None
            start_date: 开始日期，格式为'YYYYMMDD'，默认为None
            end_date: 结束日期，格式为'YYYYMMDD'，默认为None
            fields: 返回字段，默为HKConstants.HK_DAILY_ADJ_FIELDS.value

        Returns:
            DataFrame: 港股复权行情数据，包含以下字段：
            HKConstants.HK_DAILY_ADJ_FIELDS

        Api Ref:
            https://tushare.pro/document/2?doc_id=339
        """
        try:
            logger.info("开始获取港股复权行情数据")

            # 调用Tushare接口
            df = self.pro.hk_daily_adj(
                ts_code=ts_code,
                trade_date=trade_date,
                start_date=start_date,
                end_date=end_date,
                fields=fields
            )

            logger.info(f"成功获取港股复权行情数据，共{len(df)}条记录")
            return df

        except Exception as e:
            logger.error(f"获取港股复权行情数据失败: {str(e)}")
            raise


# 便捷测试函数
def get_hk_basic() -> pd.DataFrame:
    """获取港股列表的测试函数"""
    data = HKTushareData()
    return data.get_hk_basic()


def get_hk_tradecal(start_date: Optional[str] = None,
                    end_date: Optional[str] = None,
                    is_open: Optional[str] = None) -> pd.DataFrame:
    """
    获取港股交易日历的便捷函数

    Args:
        start_date: 开始日期，格式为'YYYYMMDD'
        end_date: 结束日期，格式为'YYYYMMDD'
        is_open: 是否交易，'0'表示休市，'1'表示交易

    Returns:
        DataFrame: 港股交易日历数据
    """
    data = HKTushareData()
    return data.get_hk_tradecal(start_date, end_date, is_open)


def get_hk_daily_adj(ts_code: Optional[str] = None,
                     trade_date: Optional[str] = None,
                     start_date: Optional[str] = None,
                     end_date: Optional[str] = None) -> pd.DataFrame:
    """
    获取港股复权行情数据的便捷函数

    Args:
        ts_code: 股票代码，格式为'00001.HK'
        trade_date: 交易日期，格式为'YYYYMMDD'
        start_date: 开始日期，格式为'YYYYMMDD'
        end_date: 结束日期，格式为'YYYYMMDD'

    Returns:
        DataFrame: 港股复权行情数据
    """
    data = HKTushareData()
    return data.get_hk_daily_adj(ts_code, trade_date, start_date, end_date)


def test():
    """测试"""
    data = HKTushareData()
    print("港股列表:")
    print(data.get_hk_basic())
    print("\n港股交易日历(最近30天):")
    print(data.get_hk_tradecal(start_date='20250801', end_date='20250923', is_open=HKConstants.TRADE_OPEN.value))
    print("\n港股复权数据(腾讯控股最近5天):")
    print(data.get_hk_daily_adj(ts_code='00700.HK', start_date='20250801', end_date='20250805'))


if __name__ == "__main__":
    # test()
    # print(get_hk_basic())
    print(get_hk_daily_adj(ts_code='00700.HK', start_date='20250806', end_date='20250924'))
    # print(get_hk_tradecal(start_date='20250801', end_date='20250923', is_open=HKConstants.TRADE_OPEN.value))
